import pandas as pd
import numpy as np
import tushare as ts
import json
import time
import os
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
import matplotlib.pyplot as plt
from openai import OpenAI
from tqdm.notebook import tqdm
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

token = 'd942e6ff0e981f76aaa544f84405583e0c2129a8c82213637835a099'
pro = ts.pro_api(token)

client = OpenAI(api_key="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")

def list_all_days_in_month(year_month):
    start_date = datetime.strptime(year_month, "%Y%m")
    if start_date.month == 12:
        next_month_start = datetime(start_date.year + 1, 1, 1)
    else:
        next_month_start = datetime(start_date.year, start_date.month + 1, 1)

    days = []
    current_date = start_date
    while current_date <= next_month_start:
        days.append(current_date.strftime("%Y-%m-%d"))
        current_date += timedelta(days=1)

    return days


def 获取新闻1(year_month):
    print('获取新闻短讯')
    month_dates = list_all_days_in_month(year_month)
    srcs = ['sina', 'wallstreetcn', '10jqka', 'eastmoney', 'yuncaijing']
    ddds = []
    for src in srcs:
        dds = []
        for i in tqdm(range(len(month_dates) - 1)):
            start_date = month_dates[i]
            end_date = month_dates[i + 1]
            dd = pro.news(src=src, start_date=start_date, end_date=end_date)
            dd['src'] = src
            dds.append(dd)
            ddd = pd.concat(dds)
            time.sleep(0.3)
        ddds.append(ddd)
    df = pd.concat(ddds)
    return df


def 获取新闻2(year_month):
    print('获取新闻长篇')
    month_dates = list_all_days_in_month(year_month)
    dds = []
    for i in tqdm(range(len(month_dates) - 1)):
        start_date = month_dates[i] + ' 00:00:00'
        end_date = month_dates[i + 1] + ' 00:00:00'
        dd = pro.major_news(src='', start_date=start_date,
                            end_date=end_date, fields='title,content,pub_time,src')
        dds.append(dd)
        time.sleep(0.3)
    df = pd.concat(dds)
    return df


def 获取新闻3(year_month):
    print('获取CCTV新闻')
    month_dates = list_all_days_in_month(year_month)
    dds = []
    for i in tqdm(range(len(month_dates))):
        date = month_dates[i].replace('-', '')
        dd = pro.cctv_news(date=date)
        dds.append(dd)
        time.sleep(0.3)
    df = pd.concat(dds)
    return df


def get_一个月的所有新闻(year_month):
    df1 = 获取新闻1(year_month)
    df2 = 获取新闻2(year_month)
    df3 = 获取新闻3(year_month)

    df1['tp'] = '新闻短讯'
    df2['tp'] = '新闻长篇'
    df3['tp'] = 'cctv'

    dff = pd.concat([df1, df2, df3])
    #     dff = pd.concat([df2,df3])
    #     dff = pd.concat([df2])
    return dff


def get_chatgpt_response(content):
    completion = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content":
                f"""
  # 新闻信息
  # 限制范围
  你的分析只能限于以下方面，不要做限制范围以外的分析：
  在当前的该新闻，你认为该新闻对该行业的未来走势是积极，消极 还是中性？

  # 任务
  你的任务是对我提供的行业新闻进行分析和挖掘，并对该新闻进行打分，若该新闻是积极的请给出分数:1；若是中性:0；若是消极:-1。答案中不允许其他内容！
  """},
            {"role": "user", "content": f"{content}"}
        ]
    )
    return completion.choices[0].message.content

def get_chatgpt_response2(content):
    completion = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content":
                f"""
 # 限制范围
你的分析只能限于以下方面，不要做限制范围以外的分析：
在当前的该公司的财务数据，你认为这些数据对该公司的未来走势是积极，消极 还是中性？

# 任务
你的任务是对我提供的行业基本面数据进行分析和挖掘，并对该新闻进行打分，若该新闻是积极的请给出分数:1；若是中性:0；若是消极:-1。答案中不允许其他内容！
"""
},
            {"role": "user", "content": f"{content}"}
        ]
    )
    return completion.choices[0].message.content